import asyncio
import json
import logging
import re
from typing import Dict, Any, List, Optional

# 导入您项目中的 chat.py 的功能
# 确保 review.py 和 chat.py 在同一个可以相互导入的路径下
from chat import call_llm_api, split_text

# 假设日志记录器已在文件顶部定义
logger = logging.getLogger(__name__)


def build_review_prompt(text: str, instruction: str) -> str:
    """
    构建一个指令，用于从特定文件中提取可直接应用的、具体的审核要点
    """
    return f"""你是一位资深合规审计专家。请根据"审核要求"，从"源文本"中提取可直接用于审核其他文件的**具体操作要求**，去除所有抽象原则和概括性描述。

**核心要求**:
1. **具体可操作**: 每个要点必须是可直接验证的具体要求，回答"要做什么"的问题
2. **去标识化**: 移除所有公司/部门特定名称，使用通用职能描述
3. **直接应用**: 提取的要点应能直接用于审核其他同类文件
4. **严格聚焦**: 仅提取与审核要求直接相关的要点

**你的任务**:

**审核要求**:
"{instruction}"

**源文本**:
"{text}"

**输出规范**:
1. **格式**: 每个要求为完整陈述句，用英文分号(;)分隔
2. **内容**: 
   - 必须使用"应"、"必须"、"应当"等指令性词语
   - 必须包含具体操作动作(如"制定"、"建立"、"执行")
   - 必须包含操作对象(如"采购计划"、"供应商信息")
3. **禁止**:
   - 禁止抽象原则(如"保障供给原则")
   - 禁止概括描述(如"推进信息化管理")
   - 禁止包含源文件特有标识
**示例**:
- 不合格: "企业应推进采购信息化"
- 合格: "企业应建立电子采购平台执行所有采购活动;"

**请提取可直接用于文件审核的具体操作要求:**
"""



async def _extract_review_points_from_chunk_async(text_chunk: str, instruction: str) -> str:
    """
    [内部函数] 异步处理单个文本块的审核要点提取任务。
    """
    logger.info(f"开始处理一个文本块，长度为 {len(text_chunk)}...")
    prompt = build_review_prompt(text_chunk, instruction)
    messages = [{"role": "user", "content": prompt}]

    try:
        loop = asyncio.get_running_loop()
        model_response_dict = await loop.run_in_executor(
            None,
            call_llm_api,
            messages, False  # stream=False
        )

        if not model_response_dict:
            logger.warning("一个文本块的模型调用未能返回有效内容。")
            return ""  # 返回空字符串，不影响整体流程

        model_response_text = model_response_dict.get("choices", [{}])[0].get("message", {}).get("content", "")
        return model_response_text.strip().strip('；')

    except Exception as e:
        logger.error(f"处理文本块时发生错误: {e}")
        return ""  # 同样返回空字符串


async def extract_review_points(text: str, instruction: str) -> Dict[str, Any]:
    """
    [主处理函数-重构版] 根据指令，从文本中提取审核要点，并支持长文本分片处理。
    """
    try:
        if not text or not instruction:
            return {"code": 400, "msg": "请求参数'text'和'instruction'均不能为空", "data": None}

        # 1. 分片 (Split)
        # 按照您的要求，每片25000字符
        text_chunks = split_text(text, chunk_size=25000, overlap=1000)

        # 2. 并行处理 (Map)
        logger.info(f"文本被分割为 {len(text_chunks)} 个块，开始并行提取...")
        tasks = [_extract_review_points_from_chunk_async(chunk, instruction) for chunk in text_chunks]
        chunk_results = await asyncio.gather(*tasks)

        # 3. 合并与去重 (Reduce)
        # 将所有分片的结果（逗号分隔的字符串）用逗号连接成一个更长的字符串
        # 过滤掉空字符串的结果
        full_result_string = ";".join(filter(None, chunk_results))

        # 分割成列表，并去除每个元素前后的空白
        all_points = [point.strip() for point in full_result_string.split(';') if point.strip()]

        # 使用 dict.fromkeys 来实现去重，并保持原始顺序
        unique_points = list(dict.fromkeys(all_points))

        # 重新组合成最终的逗号分隔字符串
        final_response = ";".join(unique_points)

        logger.info(f"合并去重后，最终审核要点: {final_response}")

        return {
            "code": 200,
            "msg": "操作成功",
            "data": {"points": final_response}
        }

    except Exception as e:
        logger.error(f"提取审核要点时发生严重错误: {e}", exc_info=True)
        return {
            "code": 500,
            "msg": f"服务器内部错误: {str(e)}",
            "data": None
        }


# --- 功能二: 逐条审核文本 ---

def build_batch_audit_prompt(instruction_point_batch: List[str], text: str) -> str:
    """
    构建一个指令，让模型不再生成标题。
    """
    formatted_points = "\n".join([f"- {point}" for point in instruction_point_batch])

    return f"""你是一名极其严谨细致的合规审核员。你的任务是根据一个“审核标准列表”，来判断“源始文档”中的内容是否对**列表中的每一项都**合规。

**你的思考步骤 (对列表中的每一项标准)**:
1.  **理解标准**: 首先，彻底理解“审核标准”的核心要求。
2.  **定位内容**: 在“源始文档”中，仔细查找所有与该标准主题相关的条款或描述。
3.  **逻辑分析**:
    * **合规 (通过)**: 如果文档中的条款**满足或严于**审核标准，则为“通过”。
    * **违规 (不通过)**: 如果文档中的条款与审核标准**直接冲突**，则为“不通过”。
    * **缺失 (不通过)**: 如果审核标准要求**必须存在**某个条款，但文档中完全没有提及，则为“不通过”。
4.  **构建输出**: 基于你对**所有标准**的分析，严格按照下面的JSON格式输出结果。

---
**审核标准列表 (请审核所有项)**:
{formatted_points}
---
**源始文档**:
{text}
---

**输出格式要求**:
你的回复**必须**是一个JSON对象。这个对象的**键**必须与“审核标准列表”中的审核点字符串**完全一致**。每个键对应的**值**必须是一个只包含`result`, `reason`, `suggestion`的JSON对象。
**注意**：result为通过的reason和suggestion都为null。且不存在‘部分通过’，这种情况为不通过。
**示例输出**:
```json
{{
  "年龄的要求在40岁以下即可": {{
    "result": "通过",
    "reason": null,
    "suggestion": null
  }},
  "录用的人员性别不受限制": {{
    "result": "不通过",
    "reason": "文档明确要求“录用的人员应该为男性”，与性别不受限制的标准冲突。",
    "suggestion": "请删除性别要求，以确保性别不受限制。"
  }}
}}
```
"""


async def _audit_chunk_with_batches_async(text_chunk: str, point_batches: List[List[str]]) -> List[Dict]:
    """[内部函数] 在解析时，将审核点原文(key)注入到结果字典中。"""
    tasks = []
    for batch in point_batches:
        prompt = build_batch_audit_prompt(batch, text_chunk)
        messages = [{"role": "user", "content": prompt}]
        loop = asyncio.get_running_loop()
        tasks.append(loop.run_in_executor(None, call_llm_api, messages, False))

    batch_responses = await asyncio.gather(*tasks, return_exceptions=True)

    all_audits_for_chunk = []
    for i, res in enumerate(batch_responses):
        if isinstance(res, Exception):
            logger.error(f"处理审核点批次 {i} 时API调用失败: {res}")
            continue
        if not res:
            logger.warning(f"处理审核点批次 {i} 时模型未返回内容。")
            continue

        content = res.get("choices", [{}])[0].get("message", {}).get("content", "")
        try:
            json_match = re.search(r'\{.*\}', content, re.DOTALL)
            if not json_match:
                logger.warning(f"在批次 {i} 的模型返回中未找到JSON: {content}")
                continue

            audit_results_dict = json.loads(json_match.group(0))

            # 遍历键值对，将键(审核点原文)作为'title'字段添加到值中
            for instruction_point, result_details in audit_results_dict.items():
                if isinstance(result_details, dict):
                    # 将审核点原文注入，以便后续合并逻辑使用
                    result_details['title'] = instruction_point
                    all_audits_for_chunk.append(result_details)

        except (json.JSONDecodeError, IndexError) as e:
            logger.error(f"解析批次 {i} 的JSON失败: {e}. Raw content: {content}")

    return all_audits_for_chunk


async def audit_text_by_points(text: str, instruction: str) -> Dict[str, Any]:
    """[主函数] 其余逻辑保持不变，因为它们依赖于'title'字段，而这个字段现在由我们自己提供。"""
    try:
        if not text or not instruction:
            return {"code": 400, "msg": "请求参数'text'或'instruction'不能为空", "data": None}

        instruction_points = [p.strip() for p in re.split('[,;]', instruction) if p.strip()]
        if not instruction_points:
            return {"code": 400, "msg": "从'instruction'中未能解析出任何有效的审核点", "data": None}

        batch_size = 10
        point_batches = [instruction_points[i:i + batch_size] for i in range(0, len(instruction_points), batch_size)]

        text_chunks = split_text(text, chunk_size=25000, overlap=1000)

        tasks = [_audit_chunk_with_batches_async(chunk, point_batches) for chunk in text_chunks]

        all_chunk_results = await asyncio.gather(*tasks)

        # “通过优先”合并逻辑
        final_audits = {}
        for chunk_result_list in all_chunk_results:
            for audit in chunk_result_list:
                # 这里的 'title' 现在是我们自己注入的审核点原文
                title = audit.get("title")
                if not title: continue

                if title not in final_audits or final_audits[title].get("result") == "不通过":
                    final_audits[title] = audit

        # 格式化最终输出
        final_data = []
        key_map = {
            "title": "标题", "result": "审核结果",
            "reason": "不通过原因", "suggestion": "建议描述"
        }
        for audit_item in final_audits.values():
            formatted_item = {key_map.get(k, k): v for k, v in audit_item.items()}
            final_data.append(formatted_item)

        return {"code": 200, "msg": "操作成功", "data": final_data}

    except Exception as e:
        logger.error(f"文本审核时发生严重错误: {e}", exc_info=True)
        return {"code": 500, "msg": f"服务器内部错误: {str(e)}", "data": None}


